package com.xfgg.cloud.boot.jedis;

import com.google.common.collect.Lists;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import redis.clients.jedis.*;
import redis.clients.util.JedisClusterCRC16;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 集成jedis实现批量获取
 *
 * @author xfgg
 */
public class JedisClusterPipeLine extends JedisCluster {

    volatile long refreshClusterInterval = 0;
    AtomicLong lastRefreshTimestamp = new AtomicLong();

    public JedisClusterPipeLine(Set<HostAndPort> jedisClusterNode, int connectionTimeout, int soTimeout, int maxAttempts, final GenericObjectPoolConfig poolConfig) {
        super(jedisClusterNode, connectionTimeout, soTimeout, maxAttempts, poolConfig);

        super.connectionHandler = new JedisSlotAdvancedConnectionHandler(jedisClusterNode, poolConfig, connectionTimeout, soTimeout);
    }

    public JedisSlotAdvancedConnectionHandler getConnectionHandler() {
        return (JedisSlotAdvancedConnectionHandler) this.connectionHandler;
    }

    /**
     * 刷新集群信息，当集群信息发生变更时调用
     */
    private void refreshCluster() {
        if (refreshClusterInterval <= 0) {
            connectionHandler.renewSlotCache();
            return;
        }
        long currentTimeMillis = System.currentTimeMillis();
        long last = lastRefreshTimestamp.get();
        if (currentTimeMillis - refreshClusterInterval < last) {
            return;
        }
        if (lastRefreshTimestamp.compareAndSet(last, currentTimeMillis)) {
            connectionHandler.renewSlotCache();
        }
    }

    /**
     * CRC16寻址存储数据
     *
     * @param kv     key 和 value Map
     * @param expire 过期时间
     */
    public void pipelineSet(Map<String, String> kv, Integer expire) {
        JedisSlotAdvancedConnectionHandler jedisSlotAdvancedConnectionHandler = this.getConnectionHandler();
        this.refreshCluster();
        Map<JedisPool, List<String>> poolKeys = kv.keySet().stream().collect(Collectors.groupingBy(key -> jedisSlotAdvancedConnectionHandler.getJedisPoolFromSlot(JedisClusterCRC16.getSlot(key))));
        //并行任务
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            Jedis jedis = pool.getResource();
            Pipeline pipeline = jedis.pipelined();
            List<String> keys = poolKeys.get(pool);
            for (String key : keys) {
                pipeline.setex(key, expire, kv.get(key));
            }
            pipeline.sync();
            jedis.close();
        });
    }

    /**
     * CRC16寻址获取List类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<String, Response<List<String>>> pipelineListRange(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Response<List<String>>> responses = new ConcurrentHashMap<>(allKeys.size());
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<String> keys = poolKeys.get(pool);
                for (String key : keys) {
                    responses.put(key, pipeline.lrange(key, 0, -1));
                }
                pipeline.sync();//同步提交
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取String 类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<String, Response<String>> pipelineGet(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Response<String>> responses = new ConcurrentHashMap<>(allKeys.size());
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<String> keys = poolKeys.get(pool);
                for (String key : keys) {
                    responses.put(key, pipeline.get(key));
                }
                pipeline.sync();//同步提交
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取String 类型数据
     *
     * @param allKeys key
     * @return key 直接对应的value
     */
    public Map<String, Object> pipelineValueGet(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Object> responses = new ConcurrentHashMap<>(allKeys.size());
        poolKeys.forEach((pool, keys) -> {
            Pipeline pipeline = null;
            try (Jedis jedis = pool.getResource()) {
                pipeline = jedis.pipelined();
                for (String key : keys) {
                    pipeline.get(key);
                }
                List<Object> values = pipeline.syncAndReturnAll();
                for (int i = 0; i < keys.size(); i++) {
                    if (values.get(i) == null) {
                        continue;
                    }
                    responses.put(keys.get(i), values.get(i));
                }
            } finally {
                if (pipeline != null) {
                    pipeline.clear();
                }
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取Hash类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<String, Response<Map<String, String>>> pipelineHashGetAll(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Response<Map<String, String>>> responses = new ConcurrentHashMap<>(allKeys.size());
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<String> keys = poolKeys.get(pool);
                for (String key : keys) {
                    responses.put(key, pipeline.hgetAll(key));
                }
                pipeline.sync();//同步提交
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取Set类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<String, Response<Set<String>>> pipelineSetGetAll(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Response<Set<String>>> responses = new ConcurrentHashMap<>(allKeys.size());
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<String> keys = poolKeys.get(pool);
                for (String key : keys) {
                    responses.put(key, pipeline.smembers(key));
                }
                pipeline.sync();//同步提交
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取Sort Set类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<String, Response<Set<String>>> pipelineSortSetGetAll(Collection<String> allKeys) {
        Map<JedisPool, List<String>> poolKeys = getPoolKeys(allKeys);
        Map<String, Response<Set<String>>> responses = new ConcurrentHashMap<>(allKeys.size());
        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<String> keys = poolKeys.get(pool);
                for (String key : keys) {
                    responses.put(key, pipeline.zrange(key, 0, -1));
                }
                pipeline.sync();//同步提交
            }
        });
        return responses;
    }

    /**
     * CRC16寻址获取字节类型数据
     *
     * @param allKeys key
     * @return 获取到的数据
     */
    public Map<byte[], Response<byte[]>> pipelineGetByte(Collection<byte[]> allKeys) {
        JedisSlotAdvancedConnectionHandler jedisSlotAdvancedConnectionHandler = this.getConnectionHandler();
        this.refreshCluster();
        Map<JedisPool, List<byte[]>> poolKeys = allKeys.stream().collect(Collectors.groupingBy(key -> jedisSlotAdvancedConnectionHandler.getJedisPoolFromSlot(JedisClusterCRC16.getSlot(key))));

        Map<byte[], Response<byte[]>> responses = new ConcurrentHashMap<>(allKeys.size());

        Lists.newArrayList(poolKeys.keySet()).parallelStream().forEach(pool -> {
            try (Jedis jedis = pool.getResource()) {
                Pipeline pipeline = jedis.pipelined();
                List<byte[]> keys = poolKeys.get(pool);
                for (byte[] key : keys) {
                    responses.put(key, pipeline.get(key));
                }
                pipeline.sync();
            }
        });
        return responses;
    }

    private Map<JedisPool, List<String>> getPoolKeys(Collection<String> allKeys) {
        JedisSlotAdvancedConnectionHandler jedisSlotAdvancedConnectionHandler = this.getConnectionHandler();
        this.refreshCluster();
        return allKeys.stream().collect(Collectors.groupingBy(key -> jedisSlotAdvancedConnectionHandler.getJedisPoolFromSlot(JedisClusterCRC16.getSlot(key))));
    }
}
